import tensorflow as tf
import numpy as np
import json
import glob
import colour


# 获取json文件中的计算结果
def get_R(folder_path):
    out = []
    files = glob.glob(folder_path+'\\*.json')
    for file in files:
        out.append(json.load(open(file, 'r')))
    return out


# 将结数据转化为tensorlfow格式并形成tensor slice
def res_to_data(results, test_size=0.2):
    # 分割数据集
    # 对不等间距的数据进行插值，长度200
    def data_split(x, y, test_size=0.2):
        x_train, x_test = tf.split(x, num_or_size_splits=[round((1 - test_size) * x.shape[0]),
                                                          round(test_size * x.shape[0])])
        y_train, y_test = tf.split(y, num_or_size_splits=[round((1 - test_size) * y.shape[0]),
                                                          round(test_size * y.shape[0])])
        return x_train, x_test, y_train, y_test

    params = np.array([])
    rdata = np.array([])
    # lam_fdtd = np.squeeze(np.array(results[0]['R_data'])[:, :, 0])  # 同时提取波长
    IsFirst = True
    # 循环提取分布式结算的结果，json文件
    for result in results:
        if IsFirst:
            params = np.array(result['para_array'])
            lam_fdtd = np.flip(np.array(result['R_data'])[:, :, 0], 1)
            rdata = np.flip(np.array(result['R_data'])[:, :, 1], 1)
            IsFirst = False
        else:
            params = np.concatenate([params, np.array(result['para_array'])])
            lam_fdtd = np.concatenate([lam_fdtd, np.flip(np.array(result['R_data'])[:, :, 0], 1)])
            rdata = np.concatenate([rdata, np.flip(np.array(result['R_data'])[:, :, 1], 1)])

    params = tf.cast(params, dtype=tf.float32)
    rdata = tf.cast(rdata, dtype=tf.float32)
    lam_new = np.linspace(380e-9, 780e-9, 201)
    rdata_new = np.array([np.interp(lam_new, lam_fdtd[i, :], rdata[i, :]) for i in range(lam_fdtd.shape[0])])
    x_train, x_test, y_train, y_test = data_split(params, rdata_new, test_size)

    return x_train, y_train, x_test, y_test, tf.cast(lam_new, dtype=tf.float64)


# 数据归一化及恢复
class NormPara:
    def __init__(self, x, y):
        self.min = tf.reduce_min(x, axis=0)
        self.max = tf.reduce_max(x, axis=0)

    def preprocess(self, x, y):
        # self.x = x
        # self.y = y
        # self.min = tf.reduce_min(x, axis=0)
        # self.max = tf.reduce_max(x, axis=0)
        norm = (x - self.min) / (self.max - self.min)
        return norm, y

    def process(self, x_new, y_new):
        norm_new = (x_new - self.min) / (self.max - self.min)
        return norm_new, y_new

    def back_para(self, x_t1, y_t2):
        return x_t1 * (self.max - self.min) + self.min, y_t2


# 从反射谱计算CIE1931坐标
def get_xyz_pos(spec):
    # spec: 光谱(经过插值)
    # 定义光源和角度，计算CIE1931坐标，角度2，光源D65
    # 对已经划分好的训练集和测试集进行计算，所以spec为list
    y_1, y_2 = spec[0], spec[1]
    deg, ill = colour.MSDS_CMFS['CIE 1931 2 Degree Standard Observer'], colour.SDS_ILLUMINANTS['D65']
    wl = colour.SpectralShape(380, 780, 2)
    multi_spec_train = colour.MultiSpectralDistributions(tf.transpose(y_1).numpy(), wl.range())
    multi_spec_test = colour.MultiSpectralDistributions(tf.transpose(y_2).numpy(), wl.range())  # 这个函数输入的光谱数据列为光谱
    xyz_train = colour.msds_to_XYZ(multi_spec_train, cmfs=deg, illuminant=ill, method='Integration')
    xyz_test = colour.msds_to_XYZ(multi_spec_test, cmfs=deg, illuminant=ill, method='Integration')
    xyz_train_da = tf.constant(xyz_train) / tf.reshape(tf.reduce_sum(tf.constant(xyz_train), 1), (-1, 1))
    xyz_test_da = tf.constant(xyz_test) / tf.reshape(tf.reduce_sum(tf.constant(xyz_test), 1), (-1, 1))
    return xyz_train_da, xyz_test_da


# 从反射谱计算CIE1931坐标 （速度较慢但方便求导）
def my_spec_to_xyz(spec):
    # 用于方便求导
    # 经过插值的D65光源和xyz标准刺激值
    D65 = tf.constant([49.9755, 50.91002, 51.84454, 52.77908, 53.71364, 54.6482,
                       60.26952, 65.89084, 71.51218, 77.13354, 82.7549, 84.5011,
                       86.2473, 87.99352, 89.73976, 91.486, 91.87516, 92.26432,
                       92.65348, 93.04264, 93.4318, 92.08188, 90.73196, 89.38206,
                       88.03218, 86.6823, 90.31882, 93.95534, 97.59188, 101.22844,
                       104.865, 107.2934, 109.7218, 112.1504, 114.5792, 117.008,
                       117.1688, 117.3296, 117.4904, 117.6512, 117.812, 117.2216,
                       116.6312, 116.041, 115.451, 114.861, 115.0734, 115.2858,
                       115.4982, 115.7106, 115.923, 114.5006, 113.0782, 111.6558,
                       110.2334, 108.811, 108.9194, 109.0278, 109.1364, 109.2452,
                       109.354, 109.0436, 108.7332, 108.4228, 108.1124, 107.802,
                       107.1996, 106.5972, 105.9948, 105.3924, 104.79, 105.3696,
                       105.9492, 106.529, 107.109, 107.689, 107.0322, 106.3754,
                       105.7186, 105.0618, 104.405, 104.333, 104.261, 104.1892,
                       104.1176, 104.046, 103.2368, 102.4276, 101.6184, 100.8092,
                       100., 99.26684, 98.53368, 97.80052, 97.06736, 96.3342,
                       96.22496, 96.11572, 96.00648, 95.89724, 95.788, 94.36752,
                       92.94704, 91.52656, 90.10608, 88.6856, 88.94972, 89.21384,
                       89.47796, 89.74208, 90.0062, 89.92476, 89.84332, 89.7619,
                       89.6805, 89.5991, 89.21902, 88.83894, 88.45886, 88.07878,
                       87.6987, 86.81666, 85.93462, 85.0526, 84.1706, 83.2886,
                       83.37072, 83.45284, 83.53496, 83.61708, 83.6992, 82.96472,
                       82.23024, 81.49576, 80.76128, 80.0268, 80.06436, 80.10192,
                       80.13948, 80.17704, 80.2146, 80.62724, 81.03988, 81.45252,
                       81.86516, 82.2778, 81.47908, 80.68036, 79.88164, 79.08292,
                       78.2842, 76.5716, 74.859, 73.14642, 71.43386, 69.7213,
                       70.09886, 70.47642, 70.85398, 71.23154, 71.6091, 72.15706,
                       72.70502, 73.253, 73.801, 74.349, 71.8, 69.251,
                       66.702, 64.153, 61.604, 63.26032, 64.91664, 66.57296,
                       68.22928, 69.8856, 70.92588, 71.96616, 73.00644, 74.04672,
                       75.087, 72.78812, 70.48924, 68.19038, 65.89154, 63.5927,
                       60.15778, 56.72286, 53.28796, 49.85308, 46.4182, 50.49564,
                       54.57308, 58.65052, 62.72796, 66.8054, 66.12088, 65.43636,
                       64.75184, 64.06732, 63.3828])
    s_xyz = tf.constant([[1.368000e-03, 3.900000e-05, 6.450000e-03],
                         [1.943000e-03, 5.520000e-05, 9.170000e-03],
                         [2.518000e-03, 7.140000e-05, 1.189000e-02],
                         [3.093000e-03, 8.760000e-05, 1.461000e-02],
                         [3.668000e-03, 1.038000e-04, 1.733000e-02],
                         [4.243000e-03, 1.200000e-04, 2.005000e-02],
                         [6.256400e-03, 1.752000e-04, 2.961000e-02],
                         [8.269800e-03, 2.304000e-04, 3.917000e-02],
                         [1.028320e-02, 2.856000e-04, 4.873000e-02],
                         [1.229660e-02, 3.408000e-04, 5.829000e-02],
                         [1.431000e-02, 3.960000e-04, 6.785000e-02],
                         [2.015000e-02, 5.588000e-04, 9.576000e-02],
                         [2.599000e-02, 7.216000e-04, 1.236700e-01],
                         [3.183000e-02, 8.844000e-04, 1.515800e-01],
                         [3.767000e-02, 1.047200e-03, 1.794900e-01],
                         [4.351000e-02, 1.210000e-03, 2.074000e-01],
                         [6.168400e-02, 1.768000e-03, 2.950400e-01],
                         [7.985800e-02, 2.326000e-03, 3.826800e-01],
                         [9.803200e-02, 2.884000e-03, 4.703200e-01],
                         [1.162060e-01, 3.442000e-03, 5.579600e-01],
                         [1.343800e-01, 4.000000e-03, 6.456000e-01],
                         [1.642840e-01, 5.520000e-03, 7.936000e-01],
                         [1.941880e-01, 7.040000e-03, 9.416000e-01],
                         [2.240920e-01, 8.560000e-03, 1.089600e+00],
                         [2.539960e-01, 1.008000e-02, 1.237600e+00],
                         [2.839000e-01, 1.160000e-02, 1.385600e+00],
                         [2.967760e-01, 1.388000e-02, 1.457892e+00],
                         [3.096520e-01, 1.616000e-02, 1.530184e+00],
                         [3.225280e-01, 1.844000e-02, 1.602476e+00],
                         [3.354040e-01, 2.072000e-02, 1.674768e+00],
                         [3.482800e-01, 2.300000e-02, 1.747060e+00],
                         [3.458640e-01, 2.600000e-02, 1.752070e+00],
                         [3.434480e-01, 2.900000e-02, 1.757080e+00],
                         [3.410320e-01, 3.200000e-02, 1.762090e+00],
                         [3.386160e-01, 3.500000e-02, 1.767100e+00],
                         [3.362000e-01, 3.800000e-02, 1.772110e+00],
                         [3.271200e-01, 4.240000e-02, 1.751528e+00],
                         [3.180400e-01, 4.680000e-02, 1.730946e+00],
                         [3.089600e-01, 5.120000e-02, 1.710364e+00],
                         [2.998800e-01, 5.560000e-02, 1.689782e+00],
                         [2.908000e-01, 6.000000e-02, 1.669200e+00],
                         [2.717120e-01, 6.619600e-02, 1.592888e+00],
                         [2.526240e-01, 7.239200e-02, 1.516576e+00],
                         [2.335360e-01, 7.858800e-02, 1.440264e+00],
                         [2.144480e-01, 8.478400e-02, 1.363952e+00],
                         [1.953600e-01, 9.098000e-02, 1.287640e+00],
                         [1.754160e-01, 1.005880e-01, 1.192702e+00],
                         [1.554720e-01, 1.101960e-01, 1.097764e+00],
                         [1.355280e-01, 1.198040e-01, 1.002826e+00],
                         [1.155840e-01, 1.294120e-01, 9.078880e-01],
                         [9.564000e-02, 1.390200e-01, 8.129500e-01],
                         [8.291400e-02, 1.528200e-01, 7.433960e-01],
                         [7.018800e-02, 1.666200e-01, 6.738420e-01],
                         [5.746200e-02, 1.804200e-01, 6.042880e-01],
                         [4.473600e-02, 1.942200e-01, 5.347340e-01],
                         [3.201000e-02, 2.080200e-01, 4.651800e-01],
                         [2.658800e-02, 2.310160e-01, 4.265440e-01],
                         [2.116600e-02, 2.540120e-01, 3.879080e-01],
                         [1.574400e-02, 2.770080e-01, 3.492720e-01],
                         [1.032200e-02, 3.000040e-01, 3.106360e-01],
                         [4.900000e-03, 3.230000e-01, 2.720000e-01],
                         [5.780000e-03, 3.590000e-01, 2.492400e-01],
                         [6.660000e-03, 3.950000e-01, 2.264800e-01],
                         [7.540000e-03, 4.310000e-01, 2.037200e-01],
                         [8.420000e-03, 4.670000e-01, 1.809600e-01],
                         [9.300000e-03, 5.030000e-01, 1.582000e-01],
                         [2.009400e-02, 5.444000e-01, 1.422100e-01],
                         [3.088800e-02, 5.858000e-01, 1.262200e-01],
                         [4.168200e-02, 6.272000e-01, 1.102300e-01],
                         [5.247600e-02, 6.686000e-01, 9.424000e-02],
                         [6.327000e-02, 7.100000e-01, 7.825000e-02],
                         [8.371600e-02, 7.404000e-01, 7.103200e-02],
                         [1.041620e-01, 7.708000e-01, 6.381400e-02],
                         [1.246080e-01, 8.012000e-01, 5.659600e-02],
                         [1.450540e-01, 8.316000e-01, 4.937800e-02],
                         [1.655000e-01, 8.620000e-01, 4.216000e-02],
                         [1.904800e-01, 8.804000e-01, 3.778800e-02],
                         [2.154600e-01, 8.988000e-01, 3.341600e-02],
                         [2.404400e-01, 9.172000e-01, 2.904400e-02],
                         [2.654200e-01, 9.356000e-01, 2.467200e-02],
                         [2.904000e-01, 9.540000e-01, 2.030000e-02],
                         [3.190100e-01, 9.621900e-01, 1.799000e-02],
                         [3.476200e-01, 9.703800e-01, 1.568000e-02],
                         [3.762300e-01, 9.785700e-01, 1.337000e-02],
                         [4.048400e-01, 9.867600e-01, 1.106000e-02],
                         [4.334500e-01, 9.949500e-01, 8.750000e-03],
                         [4.656600e-01, 9.949600e-01, 7.780000e-03],
                         [4.978700e-01, 9.949700e-01, 6.810000e-03],
                         [5.300800e-01, 9.949800e-01, 5.840000e-03],
                         [5.622900e-01, 9.949900e-01, 4.870000e-03],
                         [5.945000e-01, 9.950000e-01, 3.900000e-03],
                         [6.280200e-01, 9.864000e-01, 3.540000e-03],
                         [6.615400e-01, 9.778000e-01, 3.180000e-03],
                         [6.950600e-01, 9.692000e-01, 2.820000e-03],
                         [7.285800e-01, 9.606000e-01, 2.460000e-03],
                         [7.621000e-01, 9.520000e-01, 2.100000e-03],
                         [7.929400e-01, 9.356000e-01, 2.010000e-03],
                         [8.237800e-01, 9.192000e-01, 1.920000e-03],
                         [8.546200e-01, 9.028000e-01, 1.830000e-03],
                         [8.854600e-01, 8.864000e-01, 1.740000e-03],
                         [9.163000e-01, 8.700000e-01, 1.650000e-03],
                         [9.383000e-01, 8.474000e-01, 1.540000e-03],
                         [9.603000e-01, 8.248000e-01, 1.430000e-03],
                         [9.823000e-01, 8.022000e-01, 1.320000e-03],
                         [1.004300e+00, 7.796000e-01, 1.210000e-03],
                         [1.026300e+00, 7.570000e-01, 1.100000e-03],
                         [1.033480e+00, 7.318000e-01, 1.040000e-03],
                         [1.040660e+00, 7.066000e-01, 9.800000e-04],
                         [1.047840e+00, 6.814000e-01, 9.200000e-04],
                         [1.055020e+00, 6.562000e-01, 8.600000e-04],
                         [1.062200e+00, 6.310000e-01, 8.000000e-04],
                         [1.050280e+00, 6.054000e-01, 7.080000e-04],
                         [1.038360e+00, 5.798000e-01, 6.160000e-04],
                         [1.026440e+00, 5.542000e-01, 5.240000e-04],
                         [1.014520e+00, 5.286000e-01, 4.320000e-04],
                         [1.002600e+00, 5.030000e-01, 3.400000e-04],
                         [9.729700e-01, 4.786000e-01, 3.100000e-04],
                         [9.433400e-01, 4.542000e-01, 2.800000e-04],
                         [9.137100e-01, 4.298000e-01, 2.500000e-04],
                         [8.840800e-01, 4.054000e-01, 2.200000e-04],
                         [8.544500e-01, 3.810000e-01, 1.900000e-04],
                         [8.120400e-01, 3.578000e-01, 1.620000e-04],
                         [7.696300e-01, 3.346000e-01, 1.340000e-04],
                         [7.272200e-01, 3.114000e-01, 1.060000e-04],
                         [6.848100e-01, 2.882000e-01, 7.800000e-05],
                         [6.424000e-01, 2.650000e-01, 5.000000e-05],
                         [6.035000e-01, 2.470000e-01, 4.400000e-05],
                         [5.646000e-01, 2.290000e-01, 3.800000e-05],
                         [5.257000e-01, 2.110000e-01, 3.200000e-05],
                         [4.868000e-01, 1.930000e-01, 2.600000e-05],
                         [4.479000e-01, 1.750000e-01, 2.000000e-05],
                         [4.150200e-01, 1.614000e-01, 1.600000e-05],
                         [3.821400e-01, 1.478000e-01, 1.200000e-05],
                         [3.492600e-01, 1.342000e-01, 8.000000e-06],
                         [3.163800e-01, 1.206000e-01, 4.000000e-06],
                         [2.835000e-01, 1.070000e-01, 0.000000e+00],
                         [2.597800e-01, 9.780000e-02, 0.000000e+00],
                         [2.360600e-01, 8.860000e-02, 0.000000e+00],
                         [2.123400e-01, 7.940000e-02, 0.000000e+00],
                         [1.886200e-01, 7.020000e-02, 0.000000e+00],
                         [1.649000e-01, 6.100000e-02, 0.000000e+00],
                         [1.494000e-01, 5.520000e-02, 0.000000e+00],
                         [1.339000e-01, 4.940000e-02, 0.000000e+00],
                         [1.184000e-01, 4.360000e-02, 0.000000e+00],
                         [1.029000e-01, 3.780000e-02, 0.000000e+00],
                         [8.740000e-02, 3.200000e-02, 0.000000e+00],
                         [7.927400e-02, 2.900000e-02, 0.000000e+00],
                         [7.114800e-02, 2.600000e-02, 0.000000e+00],
                         [6.302200e-02, 2.300000e-02, 0.000000e+00],
                         [5.489600e-02, 2.000000e-02, 0.000000e+00],
                         [4.677000e-02, 1.700000e-02, 0.000000e+00],
                         [4.195600e-02, 1.524200e-02, 0.000000e+00],
                         [3.714200e-02, 1.348400e-02, 0.000000e+00],
                         [3.232800e-02, 1.172600e-02, 0.000000e+00],
                         [2.751400e-02, 9.968000e-03, 0.000000e+00],
                         [2.270000e-02, 8.210000e-03, 0.000000e+00],
                         [2.043180e-02, 7.388400e-03, 0.000000e+00],
                         [1.816360e-02, 6.566800e-03, 0.000000e+00],
                         [1.589540e-02, 5.745200e-03, 0.000000e+00],
                         [1.362720e-02, 4.923600e-03, 0.000000e+00],
                         [1.135900e-02, 4.102000e-03, 0.000000e+00],
                         [1.024520e-02, 3.699800e-03, 0.000000e+00],
                         [9.131400e-03, 3.297600e-03, 0.000000e+00],
                         [8.017600e-03, 2.895400e-03, 0.000000e+00],
                         [6.903800e-03, 2.493200e-03, 0.000000e+00],
                         [5.790000e-03, 2.091000e-03, 0.000000e+00],
                         [5.211800e-03, 1.882200e-03, 0.000000e+00],
                         [4.633600e-03, 1.673400e-03, 0.000000e+00],
                         [4.055400e-03, 1.464600e-03, 0.000000e+00],
                         [3.477200e-03, 1.255800e-03, 0.000000e+00],
                         [2.899000e-03, 1.047000e-03, 0.000000e+00],
                         [2.607200e-03, 9.416000e-04, 0.000000e+00],
                         [2.315400e-03, 8.362000e-04, 0.000000e+00],
                         [2.023600e-03, 7.308000e-04, 0.000000e+00],
                         [1.731800e-03, 6.254000e-04, 0.000000e+00],
                         [1.440000e-03, 5.200000e-04, 0.000000e+00],
                         [1.290000e-03, 4.658000e-04, 0.000000e+00],
                         [1.140000e-03, 4.116000e-04, 0.000000e+00],
                         [9.900000e-04, 3.574000e-04, 0.000000e+00],
                         [8.400000e-04, 3.032000e-04, 0.000000e+00],
                         [6.900000e-04, 2.490000e-04, 0.000000e+00],
                         [6.184000e-04, 2.232000e-04, 0.000000e+00],
                         [5.468000e-04, 1.974000e-04, 0.000000e+00],
                         [4.752000e-04, 1.716000e-04, 0.000000e+00],
                         [4.036000e-04, 1.458000e-04, 0.000000e+00],
                         [3.320000e-04, 1.200000e-04, 0.000000e+00],
                         [2.988000e-04, 1.080000e-04, 0.000000e+00],
                         [2.656000e-04, 9.600000e-05, 0.000000e+00],
                         [2.324000e-04, 8.400000e-05, 0.000000e+00],
                         [1.992000e-04, 7.200000e-05, 0.000000e+00],
                         [1.660000e-04, 6.000000e-05, 0.000000e+00],
                         [1.494200e-04, 5.400000e-05, 0.000000e+00],
                         [1.328400e-04, 4.800000e-05, 0.000000e+00],
                         [1.162600e-04, 4.200000e-05, 0.000000e+00],
                         [9.968000e-05, 3.600000e-05, 0.000000e+00],
                         [8.310000e-05, 3.000000e-05, 0.000000e+00],
                         [7.478000e-05, 2.700000e-05, 0.000000e+00],
                         [6.646000e-05, 2.400000e-05, 0.000000e+00],
                         [5.814000e-05, 2.100000e-05, 0.000000e+00],
                         [4.982000e-05, 1.800000e-05, 0.000000e+00],
                         [4.150000e-05, 1.500000e-05, 0.000000e+00]])
    lam = tf.linspace(380, 780, 201)
    for i in range(spec.shape[0]):
        if i == 0:
            x = tf.reshape(tf.reduce_sum(D65 * s_xyz[:, 0] * spec[i, :]), (-1, 1))
            y = tf.reshape(tf.reduce_sum(D65 * s_xyz[:, 1] * spec[i, :]), (-1, 1))
            z = tf.reshape(tf.reduce_sum(D65 * s_xyz[:, 2] * spec[i, :]), (-1, 1))
        else:
            x = tf.concat([x, tf.reshape(tf.reduce_sum(D65 * s_xyz[:, 0] * spec[i, :]), (-1, 1))], axis=0)
            y = tf.concat([y, tf.reshape(tf.reduce_sum(D65 * s_xyz[:, 1] * spec[i, :]), (-1, 1))], axis=0)
            z = tf.concat([z, tf.reshape(tf.reduce_sum(D65 * s_xyz[:, 2] * spec[i, :]), (-1, 1))], axis=0)
    pos = tf.concat([x, y, z], axis=1)
    return pos/tf.reshape(tf.reduce_sum(pos, axis=1), (-1, 1))


# 迭代处理数据，最远点采样，首先确定光谱两端，插值，找与原光谱最大距离点，采样，以此类推
def sample_spec_p(wl, spec, err_lim, iterations=50):
    # wl: 波长阵列
    # spec: 反射率
    # err_lim: 误差阈值
    # iterations: 最大迭代次数
    n_id = [0, len(wl)-1]
    p = wl.numpy()[n_id]
    err = 1e5
    # 插值
    for it in range(iterations):
        if err > err_lim:
            spec_in = np.interp(wl, p, spec.numpy()[n_id])
            dis = np.abs(np.subtract(spec_in, spec))
            err = np.mean(dis)
            n_id = np.append(n_id, np.argmax(dis))
            n_id = np.sort(n_id)
            p = wl.numpy()[n_id]
        else:
            return n_id, spec.numpy()[n_id]
    return n_id, spec.numpy()[n_id]


# gpu按需分配
def Set_GPU_Memory_Growth():
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # 设置 GPU 显存占用为按需分配
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
        except RuntimeError as e:
            # 异常处理
            print(e)
    else:
        print('No GPU')






















